//
// Created by zhangcc on 2021/5/24.
//

#include <Eigen/Core>
#include <Eigen/Dense>

#include "torchreid.h"

typedef Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> FEATURESS;
static Eigen::MatrixXf _cosine_distance(const FEATURESS & a, const FEATURESS & b, bool is_l2norm=false)
{
    FEATURESS aa = a;
    FEATURESS bb = b;
    if (is_l2norm)
    {
        for (int i = 0; i < a.rows(); ++i) {
            aa.row(i) = a.row(i) / sqrt(a.row(i).squaredNorm());
        }
        for (int i = 0; i < b.rows(); ++i) {
            bb.row(i) = b.row(i) / sqrt(b.row(i).squaredNorm());
        }
    }

    Eigen::MatrixXf res = 1. - (aa * bb.transpose()).array();

    return res;
}

static std::vector<float> nn_cosine_distance(std::vector<std::vector<float>> &feats, std::vector<std::vector<float>> &features)
{
    float * feats_ = (float *)malloc(feats.size() * feats[0].size() * sizeof(float));
    float * features_ = (float *)malloc(features.size() * features[0].size() * sizeof(float));
    int offset_bin_1=feats[0].size()*sizeof(float);
    int offset_bin_2=features[0].size()*sizeof(float);
    for (size_t i =0;i<feats.size(); i++){
        memcpy((void *)(feats_)+i*offset_bin_1, (void *)feats[i].data(), offset_bin_1);
    }
    for (size_t i =0;i<features.size(); i++){

        memcpy((void *)(features_)+i*offset_bin_2, (void *)features[i].data(), offset_bin_2);
    }
    FEATURESS track = Eigen::Map<FEATURESS>(feats_, feats.size(), feats[0].size());
    FEATURESS det = Eigen::Map<FEATURESS>(features_, features.size(), features[0].size());
    Eigen::MatrixXf distances = _cosine_distance(track, det);

    free(feats_);
    free(features_);

    Eigen::VectorXf res = distances.colwise().minCoeff().cwiseMax(0);

    return std::vector<float>{res.data(), res.data()+res.size()};
}

int main(int argc, char *argv[])
{
    std::cout << "Hello, World!" << std::endl;

    int device_id = 2;
    int max_batch_size = 2;
    cv::Size size = cv::Size(128, 256);
    TorchReID *reid = new TorchReID(device_id, max_batch_size, size);
    reid->init("../../models/osnet_x1_0-sim.onnx");

    const int batch_size = 10;
    const int max_test_iters = 100;

    cv::Mat image = cv::imread("../../images/test.jpg");
    cv::cvtColor(image, image, cv::COLOR_BGR2RGB);
    cv::resize(image, image, size);
    image.convertTo(image, CV_32FC3, 1.0f);

    std::vector<cv::Mat> inputs;
    for (size_t b = 0; b < batch_size; b++) {
        inputs.push_back(image);
    }

    struct timeval t1, t2;
    double elapsed_time;
    gettimeofday(&t1,NULL);

    std::vector <std::vector<float>> feats_track;
    std::vector <std::vector<float>> feats_det;
    for (int iter = 0; iter < max_test_iters; iter++)
    {
        feats_track = reid->extract(inputs);
        feats_det = reid->extract(inputs);
    }

    gettimeofday(&t2,NULL);
    elapsed_time = (t2.tv_sec - t1.tv_sec) + (double)(t2.tv_usec - t1.tv_usec)/1000000.0;

    for (size_t k = 0; k < feats_det[batch_size - 1].size(); k++)
    {
        std::cout << feats_det[batch_size - 1][k] << " ";
        if (k % 10 == 0)
        {
            std::cout << std::endl;
        }
    }
    std::cout << std::endl;

    std::vector<float> distance = nn_cosine_distance(feats_track, feats_det);
    for (size_t i = 0; i < distance.size(); i++)
    {
        std::cout << i << " " << distance[i] << std::endl;
    }
    std::cout << std::endl;
    std::cout << "average elapsed time = " << 1000 * elapsed_time << "ms" << std::endl;

    delete reid;

    return 0;
}

